// Copyright (c) 2025 Shenzhen Kaihong Digital Industry Development Co., Ltd.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::{collections::{hash_map::Entry, HashMap}, sync::Arc};

use lazy_static::lazy_static;
use log::error;
use maestro_derive::MsgTag;
use ohmw_derive::{Decode, Encode};
use rosrust::{api::SystemState, error::ResponseError};
use tokio::sync::Mutex;

use crate::system::SystemHandles;

#[derive(Clone, PartialEq, Eq, Encode, Decode, MsgTag)]
pub struct RosNodeState {
    pub node_name: String,
    pub publishers: Vec<String>,
    pub subscribers: Vec<String>,
    pub services: Vec<String>,
}

#[derive(Encode, Decode, MsgTag)]
#[msg_tag(case = "snake_case")]
pub struct GetRosStateRet {
    pub success: String,
    pub node_states: Vec<RosNodeState>,
}

#[derive(Debug, thiserror::Error)]
enum GetRosStateError {
    #[error("ros not inited")]
    RosNotInited,
    #[error("response error")]
    ResponseError(#[from] ResponseError),
}

lazy_static! {
    static ref LAST_STATES: Arc<Mutex<Option<HashMap<String, RosNodeState>>>> = Arc::new(Mutex::new(None));
}

impl RosNodeState {
    pub fn sort(&mut self) {
        self.publishers.sort();
        self.subscribers.sort();
        self.services.sort();
    }

    pub fn from_system_state(state: &SystemState) -> HashMap<String, RosNodeState> {
        let mut map: HashMap<String, RosNodeState> = HashMap::new();
        
        for topic in state.publishers.iter() {
            for publisher in topic.connections.iter() {
                match map.entry(publisher.to_string()) {
                    Entry::Occupied(mut occupied_entry) => {
                        occupied_entry.get_mut().publishers.push(topic.name.to_string());
                    },
                    Entry::Vacant(vacant_entry) => {
                        let node_state = RosNodeState {
                            node_name: publisher.to_string(),
                            publishers: vec![topic.name.to_string()],
                            subscribers: vec![],
                            services: vec![],
                        };
                        vacant_entry.insert(node_state);
                    },
                }
            }
        }

        for topic in state.subscribers.iter() {
            for subscriber in topic.connections.iter() {
                match map.entry(subscriber.to_string()) {
                    Entry::Occupied(mut occupied_entry) => {
                        occupied_entry.get_mut().subscribers.push(topic.name.to_string());
                    },
                    Entry::Vacant(vacant_entry) => {
                        let node_state = RosNodeState {
                            node_name: subscriber.to_string(),
                            publishers: vec![],
                            subscribers: vec![topic.name.to_string()],
                            services: vec![],
                        };
                        vacant_entry.insert(node_state);
                    },
                }
            }
        }

        for service in state.services.iter() {
            for provider in service.connections.iter() {
                match map.entry(provider.to_string()) {
                    Entry::Occupied(mut occupied_entry) => {
                        occupied_entry.get_mut().services.push(service.name.to_string());
                    },
                    Entry::Vacant(vacant_entry) => {
                        let node_state = RosNodeState {
                            node_name: provider.to_string(),
                            publishers: vec![],
                            subscribers: vec![],
                            services: vec![service.name.to_string()],
                        };
                        vacant_entry.insert(node_state);
                    },
                }
            }
        }

        map
    }
}

fn get_ros_state_impl() -> Result<HashMap<String, RosNodeState>, GetRosStateError> {
    if !SystemHandles::ensure_ros_initialized() {
        return Err(GetRosStateError::RosNotInited);
    }

    let state = match rosrust::state() {
        Ok(state) => state,
        Err(err) => {
            return Err(GetRosStateError::ResponseError(err));
        }
    };

    let mut node_states = RosNodeState::from_system_state(&state);
    
    // sort for comparison
    for node_state in node_states.values_mut() {
        node_state.sort();
    }

    Ok(node_states)
}

pub fn get_ros_state() -> GetRosStateRet {
    match get_ros_state_impl() {
        Ok(node_states) => {
            let success = "ok".to_string();
            GetRosStateRet {
                success,
                node_states: node_states.values().cloned().collect(),
            }
        },
        Err(err) => {
            error!("Failed to get system state: {}", err);
            let success = format!("getSystemState failed: {}", err);
            GetRosStateRet {
                success,
                node_states: vec![],
            }
        },
    }
}

pub async fn check_ros_state_update() -> Option<GetRosStateRet> {
    if !SystemHandles::is_ros_initialized() {
        return None;
    }

    let states = match get_ros_state_impl() {
        Ok(state) => state,
        Err(err) => {
            error!("Failed to get system state: {}", err);
            return None;
        },
    };

    let mut last_states = LAST_STATES.lock().await;
    if last_states.as_ref().is_none_or(|last_states| &states != last_states) {
        last_states.replace(states.clone());
        Some(GetRosStateRet {
            success: "ok".to_string(),
            node_states: last_states.as_mut().unwrap().values().cloned().collect(),
        })
    } else {
        None
    }
}
